package org.zebra.mybatis.plus.ext.audit.impl;

import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import lombok.Data;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.JdbcParameter;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.statement.update.UpdateSet;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.scripting.defaults.DefaultParameterHandler;
import org.zebra.mybatis.plus.ext.audit.DataAudit;
import org.zebra.mybatis.plus.ext.audit.DataUpdateLimitationException;
import org.zebra.mybatis.plus.ext.audit.OperationResult;

import java.sql.*;
import java.util.*;

/**
 * @author zhanghongbin
 */
public class UpdateDataAudit extends DataAudit {

    private Update updateStmt;
    private MappedStatement mappedStatement;
    private BoundSql boundSql;
    private int batchUpdateLimit;
    private boolean batchUpdateLimitationOpened;
    private Map<String, Integer> batchUpdateLimitMap;
    private Map<String, Set<String>> ignoredTableColumns;
    private TableInfo tableInfo;

    public UpdateDataAudit(
            Map<String, Set<String>> ignoredTableColumns,
            int batchUpdateLimit,
            boolean batchUpdateLimitationOpened,
            Map<String, Integer> batchUpdateLimitMap,
            Update updateStmt,
            MappedStatement mappedStatement,
            BoundSql boundSql,
            Connection connection,
            TableInfo tableInfo) {
        this.ignoredTableColumns = ignoredTableColumns;
        this.batchUpdateLimit = batchUpdateLimit;
        this.batchUpdateLimitationOpened = batchUpdateLimitationOpened;
        this.batchUpdateLimitMap = batchUpdateLimitMap;
        this.updateStmt = updateStmt;
        this.mappedStatement = mappedStatement;
        this.boundSql = boundSql;
        this.connection = connection;
        this.tableInfo = tableInfo;
    }

    private Connection connection;

    @Override
    public OperationResult process() {
        Expression where = updateStmt.getWhere();
        Select selectStmt = new Select();
        PlainSelect selectBody = new PlainSelect();
        Table table = updateStmt.getTable();
        selectBody.setFromItem(table);
        List<Column> updateColumns = new ArrayList<>();
        for (UpdateSet updateSet : updateStmt.getUpdateSets()) {
            updateColumns.addAll(updateSet.getColumns());
        }
        Columns2SelectItemsResult columns2SelectItems = this.buildColumns2SelectItems(updateColumns);

        selectBody.setSelectItems(columns2SelectItems.getSelectItems());
        selectBody.setWhere(where);
        selectStmt.setSelectBody(selectBody);
        BoundSql boundSql4Select = new BoundSql(
                mappedStatement.getConfiguration(),
                selectStmt.toString(),
                prepareParameterMapping4Select(boundSql.getParameterMappings(), updateStmt),
                boundSql.getParameterObject());
        PluginUtils.MPBoundSql mpBoundSql = PluginUtils.mpBoundSql(boundSql);
        Map<String, Object> additionalParameters = mpBoundSql.additionalParameters();
        if (additionalParameters != null && !additionalParameters.isEmpty()) {
            for (Map.Entry<String, Object> ety : additionalParameters.entrySet()) {
                boundSql4Select.setAdditionalParameter(ety.getKey(), ety.getValue());
            }
        }
        OperationResult result = new OperationResult();
        // 如果是逻辑删除
        if (columns2SelectItems.isLogicDelete()) {
            // 查询出所有列
            ((PlainSelect) selectStmt.getSelectBody()).setSelectItems(List.of(new AllColumns()));
            String originalData = this.buildOriginalData(
                    this.batchUpdateLimit,
                    this.batchUpdateLimitationOpened,
                    this.batchUpdateLimitMap,
                    selectStmt,
                    mappedStatement,
                    boundSql,
                    connection,
                    columns2SelectItems.getLogicDeleteColumnName());
            result.setOperation("delete");
            result.setTableName(table.getFullyQualifiedName());
            result.setRecordStatus(true);
            result.setChangedData(originalData);
            return result;
        } else {
            Map<String, Object> columnNameValMap =
                    new HashMap<>(boundSql.getParameterMappings().size());
            Map<Integer, String> columnSetIndexMap =
                    new HashMap<>(boundSql.getParameterMappings().size());
            List<Column> selectItemsFromUpdateSql = new ArrayList<>();
            int index = 0;
            for (UpdateSet updateSet : updateStmt.getUpdateSets()) {
                selectItemsFromUpdateSql.addAll(updateSet.getColumns());
                final List<Expression> updateList = updateSet.getExpressions();
                for (int i = 0; i < updateList.size(); ++i) {
                    Expression updateExps = updateList.get(i);
                    if (!(updateExps instanceof JdbcParameter)) {
                        columnNameValMap.put(
                                updateSet.getColumns().get(i).getColumnName().toUpperCase(), updateExps.toString());
                    }
                    columnSetIndexMap.put(
                            index++,
                            updateSet.getColumns().get(i).getColumnName().toUpperCase());
                }
            }
            Map<String, Object> updatedColumnDatas = this.getColumn(
                    this.tableInfo, boundSql, columnNameValMap, columnSetIndexMap, selectItemsFromUpdateSql);
            DataAudit.OriginalDataObj originalData = buildOriginalObjectData(
                    updatedColumnDatas,
                    selectStmt,
                    columns2SelectItems.getPk(),
                    mappedStatement,
                    boundSql4Select,
                    connection);
            // 更新失败
            if (originalData.isEmpty()) return null;
            result.setOperation("update");
            result.setTableName(table.getFullyQualifiedName());
            result.setRecordStatus(true);
            result.buildDataStr(this.compareAndGetUpdatedColumnDatas(
                    this.ignoredTableColumns, result.getTableName(), originalData, updatedColumnDatas));
        }
        return result;
    }

    private DataAudit.OriginalDataObj buildOriginalObjectData(
            Map<String, Object> updatedColumnDatas,
            Select selectStmt,
            Column pk,
            MappedStatement mappedStatement,
            BoundSql boundSql,
            Connection connection) {
        try (PreparedStatement statement = connection.prepareStatement(selectStmt.toString())) {
            DefaultParameterHandler parameterHandler =
                    new DefaultParameterHandler(mappedStatement, boundSql.getParameterObject(), boundSql);
            parameterHandler.setParameters(statement);
            ResultSet resultSet = statement.executeQuery();
            List<DataAudit.DataChangedRecord> originalObjectDatas = new LinkedList<>();
            int count = 0;

            while (resultSet.next()) {
                ++count;
                if (this.checkTableBatchLimitExceeded(
                        this.batchUpdateLimit,
                        this.batchUpdateLimitationOpened,
                        this.batchUpdateLimitMap,
                        selectStmt,
                        count)) {
                    logger.error("batch update limit exceed: count={}, BATCH_UPDATE_LIMIT={}", count, batchUpdateLimit);
                    throw DataUpdateLimitationException.DEFAULT;
                }
                originalObjectDatas.add(this.prepareOriginalDataObj(updatedColumnDatas, resultSet, pk));
            }
            DataAudit.OriginalDataObj result = new DataAudit.OriginalDataObj();
            result.setOriginalDataObj(originalObjectDatas);
            resultSet.close();
            return result;
        } catch (Exception e) {
            if (e instanceof DataUpdateLimitationException) {
                throw (DataUpdateLimitationException) e;
            }
            logger.error("try to get record tobe updated for selectStmt={}", selectStmt, e);
            return new DataAudit.OriginalDataObj();
        }
    }

    private DataAudit.DataChangedRecord prepareOriginalDataObj(
            Map<String, Object> updatedColumnDatas, ResultSet resultSet, Column pk) throws SQLException {
        final ResultSetMetaData metaData = resultSet.getMetaData();
        int columnCount = metaData.getColumnCount();
        List<DataAudit.DataColumnChangeResult> originalColumnDatas = new LinkedList<>();
        DataAudit.DataColumnChangeResult pkval = null;
        for (int i = 1; i <= columnCount; ++i) {
            String columnName = metaData.getColumnName(i).toUpperCase();
            DataAudit.DataColumnChangeResult col;
            Object updateVal = updatedColumnDatas.get(columnName);
            if (updateVal != null && updateVal.getClass().getCanonicalName().startsWith("java.")) {
                //   TypeReference typeReference =  (TypeReference)
                // typeHandlerRegistry.getTypeHandler(JdbcType.forCode(metaData.getColumnType(i)));
                // col = DataColumnChangeResult.constrcutByOriginalVal(columnName, resultSet.getObject(i,
                // (Class<?>)typeReference.getRawType()));
                col = DataAudit.DataColumnChangeResult.constrcutByOriginalVal(columnName, resultSet.getObject(i));
            } else {
                col = DataAudit.DataColumnChangeResult.constrcutByOriginalVal(columnName, resultSet.getObject(i));
            }
            if (pk != null && columnName.equalsIgnoreCase(pk.getColumnName())) {
                pkval = col;
            } else {
                originalColumnDatas.add(col);
            }
        }
        DataAudit.DataChangedRecord changedRecord = new DataAudit.DataChangedRecord();
        changedRecord.setOriginalColumnDatas(originalColumnDatas);
        if (pkval != null) {
            changedRecord.setPkColumnName(pkval.getColumnName());
            changedRecord.setPkColumnVal(pkval.getOriginalValue());
        }
        return changedRecord;
    }

    private Columns2SelectItemsResult buildColumns2SelectItems(List<Column> columns) {
        boolean isLogicDelete = false;
        if (columns == null || columns.isEmpty()) {
            return Columns2SelectItemsResult.build(Collections.singletonList(new AllColumns()), 0, isLogicDelete, "");
        }
        List<SelectItem> selectItems = new ArrayList<>(columns.size());
        String logicDeleteFiledName = "";
        // 获取逻辑删除字段名
        if (this.tableInfo != null) {
            logicDeleteFiledName = this.getLogicDeleteFieldName(tableInfo.getEntityType());
        }
        // 表逻辑删除列名
        String logicDeleteColumnName = "";
        for (Column column : columns) {
            if (logicDeleteFiledName.equals(StringUtils.underlineToCamel(column.getColumnName()))) {
                isLogicDelete = true;
                logicDeleteColumnName = column.getColumnName();
            }
            selectItems.add(new SelectExpressionItem(column));
        }
        if (this.tableInfo == null) {
            return Columns2SelectItemsResult.build(selectItems, 0, isLogicDelete, logicDeleteColumnName);
        }
        Column pk = new Column(this.tableInfo.getKeyColumn());
        selectItems.add(new SelectExpressionItem(pk));
        Columns2SelectItemsResult result =
                Columns2SelectItemsResult.build(selectItems, 1, isLogicDelete, logicDeleteColumnName);
        result.setPk(pk);
        return result;
    }

    //    private TableInfo getTableInfoByTableName(String tableName) {
    //        for (TableInfo tableInfo : TableInfoHelper.getTableInfos()) {
    //            if (tableName.equalsIgnoreCase(tableInfo.getTableName())) {
    //                return tableInfo;
    //            }
    //        }
    //        return null;
    //    }

    /**
     * 将update SET部分的jdbc参数去除
     *
     * @param originalMappingList 这里只会包含JdbcParameter参数
     * @param updateStmt          更新对象
     * @return 列表
     */
    private List<ParameterMapping> prepareParameterMapping4Select(
            List<ParameterMapping> originalMappingList, Update updateStmt) {
        List<Expression> updateValueExpressions = new ArrayList<>();
        for (UpdateSet updateSet : updateStmt.getUpdateSets()) {
            updateValueExpressions.addAll(updateSet.getExpressions());
        }
        int removeParamCount = 0;
        for (Expression expression : updateValueExpressions) {
            if (expression instanceof JdbcParameter) {
                ++removeParamCount;
            }
        }
        return originalMappingList.subList(removeParamCount, originalMappingList.size());
    }

    @Data
    public static class Columns2SelectItemsResult {

        private Column pk;
        /**
         * all column with additional columns: ID, etc.
         */
        private List<SelectItem> selectItems;
        /**
         * newly added column count from meta data.
         */
        private int additionalItemCount;

        private boolean isLogicDelete = false;

        private String logicDeleteColumnName = "";

        public static Columns2SelectItemsResult build(
                List<SelectItem> selectItems,
                int additionalItemCount,
                boolean isLogicDelete,
                String logicDeleteColumnName) {
            Columns2SelectItemsResult result = new Columns2SelectItemsResult();
            result.setSelectItems(selectItems);
            result.setAdditionalItemCount(additionalItemCount);
            result.setLogicDelete(isLogicDelete);
            result.setLogicDeleteColumnName(logicDeleteColumnName);
            return result;
        }
    }
}
